package main

import (
	"io"
	"log"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/exec"
    _ "embed"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"
)

//go:embed assets/terminal.html
var htmlContent string

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 仅测试用！生产环境请严格校验 Origin
	},
}
func terminalHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket 升级失败:", err)
		return
	}
	defer conn.Close()

	cmd := exec.Command("zsh")
	cmd.Env = os.Environ()

	ptmx, err := pty.Start(cmd)
	if err != nil {
		log.Println("启动 PTY 失败:", err)
		return
	}
	defer ptmx.Close()


	// 从 WebSocket → PTY
	go func() {
		defer ptmx.Close()
		for {
			_, data, err := conn.ReadMessage()
			if err != nil {
				return
			}
			ptmx.Write(data)
		}
	}()

	// 从 PTY → WebSocket
	go func() {
		defer conn.Close()
		buf := make([]byte, 1024)
		for {
			n, err := ptmx.Read(buf)
			if err != nil {
				if err != io.EOF {
					log.Println("PTY 读取错误:", err)
				}
				return
			}
			if n > 0 {
			    err = conn.WriteMessage(websocket.TextMessage, buf[:n])
				if err != nil {
					return
				}
			}
		}
	}()

	cmd.Wait()
}
func main() {
	port := flag.Int("port", 54321, "服务监听端口")


	http.HandleFunc("/ws", terminalHandler)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write([]byte(htmlContent))
	})

	addr := fmt.Sprintf(":%d", *port)
	url := fmt.Sprintf("http://localhost:%d", *port)
	log.Printf("🚀 Web 终端启动中: %s", url)

	log.Fatal(http.ListenAndServe(addr, nil))

}